#!/usr/bin/env python3
# Licensed under the MIT license.

"""
rStar多模态推理判别器 - 基于推理一致性的候选答案筛选系统

本模块实现了多模态推理结果的判别和筛选功能，通过推理一致性检查来选择最优答案。
该系统扩展了原有的判别框架，支持包含图像信息的推理结果筛选。

核心功能:
1. 多模态候选答案加载：从生成结果中提取候选答案
2. 推理一致性检查：验证推理过程的逻辑一致性
3. 答案筛选和排序：基于多种策略选择最优答案
4. 批量处理：支持大规模多模态推理结果的筛选

使用方法:
    python run_src/do_multimodal_discriminate.py \
        --config configs/multimodal_config.yaml \
        --results_dir outputs/multimodal_results \
        --output_dir outputs/discriminated_results
"""

import sys
import os
import json
from tqdm import tqdm
from collections import defaultdict
from copy import deepcopy
from datetime import datetime
from pathlib import Path

sys.path.append(".")

from common.config_loader import load_multimodal_config
from common.utils import fix_seeds, read_json, read_txt
from eval_src.Evaluator import GSM8KEvaluator
from run_src.rstar_utils import concat_solution_trace, mask_solution_trace
from models.multimodal_model import MultimodalModelInterface
import logging


class MultimodalCandidate:
    """
    多模态候选答案类
    
    该类扩展了原有的Candidate类，增加了图像信息的支持。
    """
    
    def __init__(
        self,
        solution_trace,
        masked_solution_trace_list,
        final_step,
        final_answer,
        candidate_id,
        image_path=None,
        freq=1,
        trace_reward=1.0,
        c_type="default",
    ):
        """
        初始化多模态候选答案
        
        Args:
            solution_trace: 解题过程轨迹
            masked_solution_trace_list: 掩码后的解题过程列表
            final_step: 最终步骤
            final_answer: 最终答案
            candidate_id: 候选答案ID
            image_path: 图像路径（新增）
            freq: 频率
            trace_reward: 轨迹奖励
            c_type: 候选类型
        """
        self.solution_trace = solution_trace
        self.masked_solution_trace_list = masked_solution_trace_list
        self.final_step = final_step
        self.final_answer = final_answer
        self.id = candidate_id
        self.image_path = image_path  # 新增图像路径
        self.freq = freq
        self.trace_reward = trace_reward
        self.c_type = c_type

    def __str__(self):
        image_info = f" (图像: {self.image_path})" if self.image_path else ""
        return f"Candidate {self.id}: {self.final_answer}{image_info}"

    def to_dict(self):
        """转换为字典格式"""
        return {
            "solution_trace": self.solution_trace,
            "masked_solution_trace_list": self.masked_solution_trace_list,
            "final_step": self.final_step,
            "final_answer": self.final_answer,
            "id": self.id,
            "image_path": self.image_path,
            "freq": self.freq,
            "trace_reward": self.trace_reward,
            "c_type": self.c_type
        }

    @classmethod
    def from_dict(cls, data):
        """从字典创建候选答案"""
        return cls(
            solution_trace=data["solution_trace"],
            masked_solution_trace_list=data["masked_solution_trace_list"],
            final_step=data["final_step"],
            final_answer=data["final_answer"],
            candidate_id=data["id"],
            image_path=data.get("image_path"),
            freq=data.get("freq", 1),
            trace_reward=data.get("trace_reward", 1.0),
            c_type=data.get("c_type", "default")
        )


class MultimodalDiscriminator:
    """
    多模态推理判别器
    
    该类扩展了原有的判别器功能，支持多模态推理结果的筛选。
    """
    
    def __init__(self, config, evaluator):
        """
        初始化多模态判别器
        
        Args:
            config: 配置对象
            evaluator: 评估器实例
        """
        self.config = config
        self.evaluator = evaluator
        
        # 初始化多模态模型接口
        self.model_interface = MultimodalModelInterface(config)
        
        # 加载提示模板（复用原有逻辑）
        self._load_prompts()
        
        # 设置判别参数
        self.threshold = 0.999
        self.mask_left_boundary = 0.2
        self.mask_right_boundary = 0.5
        self.num_masked_solution_traces = 4
        self.rc_mode = "mid"
        self.rc_temperature = 1.0
        self.rc_n_completions = 1
        self.rc_criteria = "reward"
        
        # 设置日志
        self.logger = logging.getLogger(__name__)
    
    def _load_prompts(self):
        """加载提示模板"""
        # 使用GSM8K的提示模板作为默认
        dataset_name = "GSM8K"
        
        fewshot_config_path = f"prompts/{dataset_name}/fewshot_cot/fewshot_cot_config.json"
        fewshot_prompt_path = f"prompts/{dataset_name}/fewshot_cot/fewshot_cot_prompt.txt"
        
        if os.path.exists(fewshot_config_path):
            self.fewshot_config = read_json(fewshot_config_path)
            self.fewshot_template = self.fewshot_config["prompt_template"]
            self.stop_tokens = self.fewshot_config["stop_tokens"]
        else:
            # 默认配置
            self.fewshot_template = "{examples}\n\n{instruction}\n"
            self.stop_tokens = ["Q:", "Question:", "\n\n"]
        
        if os.path.exists(fewshot_prompt_path):
            self.fewshot_prompt = read_txt(fewshot_prompt_path)
        else:
            self.fewshot_prompt = ""
    
    def _filter_none(self, candidates):
        """过滤空答案"""
        filtered = [c for c in candidates if c.final_answer is not None]
        self.logger.info(f"过滤空答案: {len(candidates)} -> {len(filtered)}")
        return filtered

    def _filter_long(self, candidates):
        """过滤过长答案"""
        filtered = [c for c in candidates if len(str(c.final_answer)) <= 100]
        self.logger.info(f"过滤长答案: {len(candidates)} -> {len(filtered)}")
        return filtered

    def _filter_reasoning_consistency(self, problem, candidates, image_path=None, aux={}):
        """
        基于推理一致性过滤候选答案
        
        Args:
            problem: 问题文本
            candidates: 候选答案列表
            image_path: 图像路径
            aux: 辅助信息
            
        Returns:
            过滤后的候选答案列表
        """
        if not candidates:
            return []
        
        problem_id = aux.get("problem_id", "unknown")
        
        # 准备生成输入
        gen_input_list = []
        ground_truth_list = []
        c_completion_num_list = []
        
        for c in candidates:
            if c.c_type == "default":
                if not hasattr(c, 'masked_solution_trace_list') or not c.masked_solution_trace_list:
                    # 如果没有掩码轨迹，创建默认的
                    c.masked_solution_trace_list = [c.solution_trace[:len(c.solution_trace)//2]]
                
                for masked_solution_trace in c.masked_solution_trace_list:
                    for _ in range(self.rc_n_completions):
                        prompt = self.fewshot_template.format(
                            examples=self.fewshot_prompt, 
                            instruction=problem
                        ) + masked_solution_trace
                        gen_input_list.append(prompt)
                        ground_truth_list.append(c.final_answer)
                
                c_completion_num_list.append(len(c.masked_solution_trace_list) * self.rc_n_completions)
        
        if not gen_input_list:
            self.logger.warning("没有生成输入，返回原始候选")
            return candidates
        
        # 准备图像路径列表
        image_paths = [image_path] * len(gen_input_list) if image_path else None
        
        # 使用多模态模型生成
        try:
            messages = [{'content': prompt} for prompt in gen_input_list]
            gen_output_list = self.model_interface.generate(messages, image_paths)
        except Exception as e:
            self.logger.error(f"多模态生成失败: {e}")
            return candidates
        
        # 保存生成结果
        results_data = {
            "problem_id": problem_id,
            "image_path": image_path,
            "gen_output_list": gen_output_list,
            "candidates": [c.to_dict() for c in candidates]
        }
        
        # 分析一致性并过滤
        filtered_candidates = self._analyze_consistency(
            candidates, gen_output_list, c_completion_num_list
        )
        
        self.logger.info(f"推理一致性过滤: {len(candidates)} -> {len(filtered_candidates)}")
        return filtered_candidates
    
    def _analyze_consistency(self, candidates, gen_output_list, c_completion_num_list):
        """
        分析推理一致性
        
        Args:
            candidates: 原始候选列表
            gen_output_list: 生成输出列表
            c_completion_num_list: 每个候选的完成数量列表
            
        Returns:
            过滤后的候选列表
        """
        if not gen_output_list:
            return candidates
        
        # 计算每个候选的一致性分数
        filtered_candidates = []
        output_idx = 0
        
        for i, candidate in enumerate(candidates):
            if i < len(c_completion_num_list):
                completion_num = c_completion_num_list[i]
                candidate_outputs = gen_output_list[output_idx:output_idx + completion_num]
                output_idx += completion_num
                
                # 计算一致性分数
                consistency_score = self._calculate_consistency_score(
                    candidate, candidate_outputs
                )
                
                # 根据阈值过滤
                if consistency_score >= self.threshold * 0.5:  # 降低阈值以保留更多候选
                    filtered_candidates.append(candidate)
        
        return filtered_candidates if filtered_candidates else candidates[:1]  # 至少保留一个
    
    def _calculate_consistency_score(self, candidate, outputs):
        """
        计算候选答案的一致性分数
        
        Args:
            candidate: 候选答案
            outputs: 生成输出列表
            
        Returns:
            一致性分数 (0-1)
        """
        if not outputs:
            return 0.0
        
        # 提取答案并计算一致性
        extracted_answers = []
        for output in outputs:
            try:
                extracted_answer = self.evaluator.extract_answer_from_model_completion(output)
                if extracted_answer:
                    extracted_answers.append(extracted_answer)
            except:
                continue
        
        if not extracted_answers:
            return 0.0
        
        # 计算与原答案的一致性
        consistent_count = 0
        for extracted_answer in extracted_answers:
            if self.evaluator.check_answers_equiv(extracted_answer, candidate.final_answer):
                consistent_count += 1
        
        return consistent_count / len(extracted_answers)
    
    def select_best_candidate(self, problem, candidates, gt_answer=None, image_path=None, aux={}):
        """
        选择最佳候选答案
        
        Args:
            problem: 问题文本
            candidates: 候选答案列表
            gt_answer: 标准答案
            image_path: 图像路径
            aux: 辅助信息
            
        Returns:
            最佳候选答案
        """
        self.logger.info(f"开始选择最佳候选，候选数量: {len(candidates)}")
        self.logger.info(f"标准答案: {gt_answer}")
        
        if not candidates:
            return None
        
        # 预过滤
        unfiltered_candidates = candidates
        prefiltered_candidates = self._filter_none(candidates)
        prefiltered_candidates = self._filter_long(prefiltered_candidates)
        
        self.logger.info(f"预过滤后候选: {[c.final_answer for c in prefiltered_candidates]}")
        
        # 推理一致性过滤
        filtered_candidates = self._filter_reasoning_consistency(
            problem, prefiltered_candidates, image_path, aux
        )
        
        self.logger.info(f"一致性过滤后候选: {[c.final_answer for c in filtered_candidates]}")
        
        # 选择最佳候选
        winner = self._find_winner(prefiltered_candidates, filtered_candidates, gt_answer)
        
        return winner
    
    def _find_winner(self, prefiltered_candidates, filtered_candidates, gt_answer):
        """
        从过滤后的候选中选择获胜者
        
        Args:
            prefiltered_candidates: 预过滤候选
            filtered_candidates: 过滤后候选
            gt_answer: 标准答案
            
        Returns:
            获胜候选
        """
        if not filtered_candidates:
            # 如果没有过滤后的候选，从预过滤中选择
            if prefiltered_candidates:
                return prefiltered_candidates[0]
            return None
        
        # 计算每个答案的分数
        answer_scores = defaultdict(float)
        answer_candidates = defaultdict(list)
        
        for candidate in filtered_candidates:
            answer = candidate.final_answer
            answer_candidates[answer].append(candidate)
            answer_scores[answer] += candidate.trace_reward
        
        # 选择分数最高的答案
        if answer_scores:
            best_answer = max(answer_scores.keys(), key=lambda x: answer_scores[x])
            return answer_candidates[best_answer][0]
        
        return filtered_candidates[0]


def load_candidates_from_results(results_path, image_path=None):
    """
    从结果文件加载候选答案
    
    Args:
        results_path: 结果文件路径
        image_path: 图像路径
        
    Returns:
        候选答案列表
    """
    candidates = []
    
    try:
        with open(results_path, 'r', encoding='utf-8') as f:
            for line_idx, line in enumerate(f):
                try:
                    result = json.loads(line.strip())
                    
                    # 从搜索结果中提取候选
                    search_result = result.get('search_result', {})
                    model_solutions = search_result.get('model_solutions', [])
                    
                    if model_solutions:
                        for i, solution in enumerate(model_solutions):
                            # 创建候选答案
                            candidate = MultimodalCandidate(
                                solution_trace=solution,
                                masked_solution_trace_list=[solution[:len(solution)//2]],
                                final_step=solution,
                                final_answer=result.get('model_answer', ''),
                                candidate_id=f"{line_idx}_{i}",
                                image_path=image_path
                            )
                            candidates.append(candidate)
                    else:
                        # 如果没有详细解题过程，使用基本信息
                        candidate = MultimodalCandidate(
                            solution_trace=result.get('model_completion', ''),
                            masked_solution_trace_list=[''],
                            final_step=result.get('model_completion', ''),
                            final_answer=result.get('model_answer', ''),
                            candidate_id=str(line_idx),
                            image_path=image_path
                        )
                        candidates.append(candidate)
                        
                except json.JSONDecodeError:
                    continue
                    
    except FileNotFoundError:
        print(f"结果文件不存在: {results_path}")
        return []
    
    return candidates


def main():
    """主函数"""
    from argparse import ArgumentParser
    
    parser = ArgumentParser(description="多模态推理判别器")
    
    # 基础配置
    parser.add_argument("--config", type=str, required=True, help="配置文件路径")
    parser.add_argument("--results_dir", type=str, required=True, help="推理结果目录")
    parser.add_argument("--seed", type=int, default=42, help="随机种子")
    parser.add_argument("--note", type=str, default="multimodal_discriminate", help="实验标签")
    
    # 输出配置
    parser.add_argument("--output_dir", type=str, help="输出目录")
    
    # 处理范围
    parser.add_argument("--start_idx", type=int, default=-1, help="开始索引")
    parser.add_argument("--end_idx", type=int, default=-1, help="结束索引")
    
    # 判别参数
    parser.add_argument("--threshold", type=float, default=0.999, help="判别阈值")
    
    args = parser.parse_args()
    
    # 设置随机种子
    fix_seeds(args.seed)
    
    # 设置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    logger = logging.getLogger(__name__)
    
    # 加载配置
    logger.info(f"加载配置文件: {args.config}")
    config = load_multimodal_config(args.config)
    
    # 设置输出目录
    if args.output_dir:
        output_dir = args.output_dir
    else:
        timestamp = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
        output_dir = os.path.join(args.results_dir, f"discriminated_{timestamp}---{args.note}")
    
    os.makedirs(output_dir, exist_ok=True)
    
    # 初始化评估器和判别器
    evaluator = GSM8KEvaluator()
    discriminator = MultimodalDiscriminator(config, evaluator)
    
    # 查找结果文件
    results_path = os.path.join(args.results_dir, "results.jsonl")
    if not os.path.exists(results_path):
        logger.error(f"结果文件不存在: {results_path}")
        return
    
    logger.info(f"开始处理结果文件: {results_path}")
    
    # 统计变量
    num_correct = 0
    num_correct_majvote = 0
    num_correct_limit = 0
    num_tested = 0
    total_candidates = 0
    
    # 处理每个样本
    results = []
    
    try:
        with open(results_path, 'r', encoding='utf-8') as f:
            lines = f.readlines()
            
        if args.start_idx >= 0 and args.end_idx >= 0:
            lines = lines[args.start_idx:args.end_idx]
        
        for line_idx, line in enumerate(tqdm(lines, desc="处理判别任务")):
            try:
                result = json.loads(line.strip())
                
                # 提取基本信息
                problem = result.get('question', '')
                gt_answer = result.get('gold_answer', '')
                image_path = result.get('image_path')
                
                # 加载候选答案
                candidates = [MultimodalCandidate(
                    solution_trace=result.get('model_completion', ''),
                    masked_solution_trace_list=[''],
                    final_step=result.get('model_completion', ''),
                    final_answer=result.get('model_answer', ''),
                    candidate_id=f"{line_idx}_0",
                    image_path=image_path
                )]
                
                total_candidates += len(candidates)
                
                # 选择最佳候选
                winner = discriminator.select_best_candidate(
                    problem=problem,
                    candidates=candidates,
                    gt_answer=gt_answer,
                    image_path=image_path,
                    aux={"problem_id": line_idx}
                )
                
                # 评估结果
                if winner and gt_answer:
                    predicted_answer = winner.final_answer
                    
                    # 检查正确性
                    correct = evaluator.check_answers_equiv(predicted_answer, gt_answer)
                    correct_majvote = correct  # 简化处理
                    correct_limit = correct
                    
                    num_correct += int(correct)
                    num_correct_majvote += int(correct_majvote)
                    num_correct_limit += int(correct_limit)
                
                num_tested += 1
                
                # 保存结果
                result_data = {
                    "problem_id": line_idx,
                    "question": problem,
                    "image_path": image_path,
                    "gold_answer": gt_answer,
                    "predicted_answer": winner.final_answer if winner else None,
                    "correct": correct if winner and gt_answer else False,
                    "num_candidates": len(candidates),
                    "winner": winner.to_dict() if winner else None
                }
                results.append(result_data)
                
            except json.JSONDecodeError:
                continue
                
    except Exception as e:
        logger.error(f"处理过程中出错: {e}")
        return
    
    # 保存结果
    output_path = os.path.join(output_dir, "discriminated_results.jsonl")
    with open(output_path, 'w', encoding='utf-8') as f:
        for result in results:
            f.write(json.dumps(result, ensure_ascii=False) + '\n')
    
    # 保存统计信息
    stats = {
        "num_correct": num_correct,
        "num_correct_majvote": num_correct_majvote,
        "num_correct_limit": num_correct_limit,
        "num_tested": num_tested,
        "accuracy": num_correct / num_tested if num_tested > 0 else 0,
        "majority_vote_accuracy": num_correct_majvote / num_tested if num_tested > 0 else 0,
        "limit_accuracy": num_correct_limit / num_tested if num_tested > 0 else 0,
        "avg_num_candidates": total_candidates / num_tested if num_tested > 0 else 0,
        "config": {
            "threshold": args.threshold,
            "model_type": config.model_type
        }
    }
    
    stats_path = os.path.join(output_dir, "stats.json")
    with open(stats_path, 'w', encoding='utf-8') as f:
        json.dump(stats, f, ensure_ascii=False, indent=2)
    
    # 打印结果
    print(f"\n{'='*60}")
    print("多模态推理判别完成")
    print(f"{'='*60}")
    print(f"处理样本数: {num_tested}")
    print(f"正确答案数: {num_correct}")
    print(f"准确率: {num_correct/num_tested:.4f}" if num_tested > 0 else "准确率: N/A")
    print(f"多数投票准确率: {num_correct_majvote/num_tested:.4f}" if num_tested > 0 else "多数投票准确率: N/A")
    print(f"限制准确率: {num_correct_limit/num_tested:.4f}" if num_tested > 0 else "限制准确率: N/A")
    print(f"平均候选数: {total_candidates/num_tested:.2f}" if num_tested > 0 else "平均候选数: N/A")
    print(f"结果保存到: {output_path}")
    print(f"统计信息保存到: {stats_path}")
    print(f"{'='*60}")


if __name__ == "__main__":
    main() 